# u^(4)=f 真解 u=x^4*(1-x)^4
# BCs: u(0)=u(1)=u'(1)=u'(0)=0
# 分解为两个方程 
# -u''=v 两边乘q
# -v''=f 两边乘p
# Weak Form: (v', p') + (u', q') - (v, q) = (f, p)

# 方程边界处理
# 首先说明f'(x)=0和f''(x)=0之间没有任何关系，两者不能互推！
# 因此把方程拆分之后 v(0)=v(1)=0用来描述边界条件是错误的
# 所以这个方法不太能解决这个问题
# poisson4b程序给出了一种可行的解决方案
try:
    from petsc4py import PETSc
    import dolfinx
    if not dolfinx.has_petsc:
        print("This demo requires DOLFINx to be compiled with PETSc enabled.")
        exit(0)
        
except ModuleNotFoundError:
    print("This demo requires petsc4py.")
    exit(0)

from mpi4py import MPI

import numpy as np
from petsc4py.PETSc import ScalarType
from basix.ufl import element, mixed_element
from dolfinx import default_real_type, fem, io, mesh
from dolfinx.fem.petsc import LinearProblem
from ufl import Measure, SpatialCoordinate, TestFunctions, TrialFunctions, inner, grad, dx


def true_u(x):
    return x[0]**4 * (1 - x[0])**4

n_list = [
    4, 8, 16, 32, 64
]
idx = int(input("Please input the index of n_list: "))
n = n_list[idx]
msh = mesh.create_unit_interval(
    MPI.COMM_WORLD, n, )

k = 2
Q_el = element("P", msh.basix_cell(), k, dtype=default_real_type)
P_el = element("P", msh.basix_cell(), k, dtype=default_real_type)
V_el = mixed_element([Q_el, P_el])
V = fem.functionspace(msh, V_el)

(v, u) = TrialFunctions(V)
(p, q) = TestFunctions(V)

x = SpatialCoordinate(msh)
f = 1680*x[0]**4 - 3360*x[0]**3 + 2160*x[0]**2 - 480*x[0] + 24

# 注意不要写反顺序 u对应q v对应p
a = inner(grad(v), grad(p)) * dx + inner(grad(u), grad(q)) * dx - inner(v, q) * dx
L = inner(f, p) * dx

# Get subspace of V
V0, V1 = V.sub(0), V.sub(1)
Q, _ = V0.collapse()
S, _ = V1.collapse()
fdim = msh.topology.dim - 1
facets = mesh.locate_entities_boundary(
    msh, fdim, lambda x: np.isclose(x[0], 1.0) | np.isclose(x[0], 0.0))

dofs = fem.locate_dofs_topological(V=V0, entity_dim=0, entities=facets)

bc_u = fem.dirichletbc(ScalarType(0), dofs, V1)
bc_v = fem.dirichletbc(ScalarType(0), dofs, V0)
bcs = [bc_u, bc_v]

problem = LinearProblem(
    a,
    L,
    bcs=bcs,
    petsc_options={
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "superlu_dist",
    },
)
try:
    w_h = problem.solve()
except PETSc.Error as e:  # type: ignore
    if e.ierr == 92:
        print("The required PETSc solver/preconditioner is not available. Exiting.")
        print(e)
        exit(0)
    else:
        raise e

# 分离视图 两个变量的值是一样的 不清楚原因 都是混合空间
v_h, u_h = w_h.split() 

# 创建真实解的函数对象
u_true = fem.Function(S)
u_true.interpolate(true_u)

# 计算误差的表达式
error_expr = u_true - u_h
l2_error = fem.assemble_scalar(fem.form(inner(error_expr, error_expr) * dx))
l2_error = np.sqrt(l2_error)

# 计算 H1 范数误差
h1_error_expr = inner(grad(error_expr), grad(error_expr)) * \
    dx + inner(error_expr, error_expr) * dx
h1_error = fem.assemble_scalar(fem.form(h1_error_expr))
h1_error = np.sqrt(h1_error)

# Linf误差
# u_error = fem.Function(S)
# u_error.x.array[:] = u_true.x.array - u_h.x.array[2*n+1:]
# linf_error = dolfinx.la.norm(u_error.x, dolfinx.la.Norm.linf)

print(f"N = {n}.")
print(f"L2 norm error: {l2_error}")
print(f"H1 norm error: {h1_error}")
# print(f"Linf norm error: {linf_error}")
